from llama_cpp import Llama


if __name__ == '__main__':

    # 指定模型路径，模型文件应该是.gguf格式的
    model_path = r"D:\ai\model\causallm_7b-dpo-alpha.Q5_K_M.gguf"

    llm = Llama(
        model_path=model_path,
        n_ctx=2048,
        chat_format="llama-2"
    )

    user_msg = """提取下面这段文本的实体及其属性信息：035型潜艇的先进推进系统是其强大性能的核心所在。通过先进的技术和工艺，潜艇可以在水下保持高效稳定的航行状态，实现长时间的连续航行。这一特点在实际应用中发挥着重要作用，比如在海上巡逻和情报搜集中，
潜艇需要长时间潜伏在水下执行任务，而035型潜艇的续航能力为其提供了充分的支持。，
除了优秀的水下航行能力，035型潜艇还具备多种功能，使其成为海军作战中的重要力量。例如，该型潜艇可以搭载鱼雷、反舰导弹等多种武器装备，实施海上打击和防御任务。
同时，它还具备布设水雷、侦察监视等功能，为海上作战提供了多样化的选择，增强了海军的整体作战能力。
在复杂的海况下，035型潜艇展现出了出色的适应能力。无论是在恶劣天气下还是在深海中，潜艇都能保持稳定的航行状态，确保舰员和装备的安全。这为海军在恶劣环境下执行任务提供了有力支持，保障了海上安全和国家利益。
中国自行研制的035型潜艇不仅仅是一种装备，更代表着国家在军事技术领域的创新和发展。通过自主研发和不断改进，我国海军的潜艇部队逐渐实现了由跟随到领先的转变，为维护国家海洋权益和安全作出了重要贡献。
随着中国自行研制的035型潜艇逐渐走向国际舞台，其性能和技术也受到了世界各国的关注和认可。这种国产潜艇的出口意味着中国海军装备水平的提升和国际地位的提高，对于推动国际军事技术合作和交流具有积极的意义"""
    prompt = 'Human:'+user_msg+'\nAssistant:'
    response = llm(prompt, stop=['Human:'])
    print(response)